from osgeo import gdal,osr,gdalconst,ogr
from xml.dom.minidom import parse
# import datetime as dt
import json
import daemon
import redis
import psycopg2
import shutil
import zipfile
# from rocketmq.client import Producer, Message
import os
# import rey
import time


def name_2_count(cargs,name):
    connection = psycopg2.connect(**cargs)
    connection.autocommit = True
    cursor = connection.cursor()
    query = f"""select tif_path from sentinel2_l1c where name = '{name}' and tif_path is not null;"""
    cursor.execute(query)
    data = cursor.fetchone()
    cursor.close()
    connection.close()
    return data


def read_xmlfile(xml_path,select_node):
    domTree = parse(xml_path)
    collection = domTree.documentElement
    select_element = [collection]
    for node_dict in select_node:
        tag_name = node_dict["tag_name"]
        if node_dict["attribute"] is None:
            select_element = [
                element_sub
                for element in select_element
                for element_sub in element.getElementsByTagName(tag_name)
            ]
        else:
            attribute_name = node_dict["attribute"]["name"]
            attribute_value = node_dict["attribute"]["value"]
            select_element = [
                element_sub
                for element in select_element
                for element_sub in element.getElementsByTagName(tag_name)
                if element_sub.getAttribute(attribute_name) == attribute_value
            ]
    datas = [
        element.firstChild.data
        for element in select_element
    ]
    return datas


def read_cloud(zip_path):
    select_node = [
        {
            "tag_name": "n1:Quality_Indicators_Info",
            "attribute": None
        },
        {
            "tag_name": "Cloud_Coverage_Assessment",
            "attribute": None
        }
    ]
    xml_namere = "^S2.*\.SAFE/MTD_MSIL1C\.xml$"
    zip_path = zip_path.replace('"','')
    zip_file = zipfile.ZipFile(zip_path, "r")
    cloud = "00"
    for file_info in zip_file.infolist():
        if re.match(xml_namere, file_info.filename):
            xml_path = zip_file.open(file_info.filename, "r")
            datas = read_xmlfile(xml_path, select_node)
            cloud = round(float(datas[0]))
            cloud = f"0{cloud}" if len(f"{cloud}") == 1 else f"{cloud}"
    return float(cloud)


# 发送数据到消息队列
def send_data(json_dir,zip_path,tif_path,redisr,topic="henan_rsanalysis_s2_1c_collect",addr='192.168.1.70:9876'):
    name = os.path.basename(zip_path).replace(".zip","")
    tile =  name.split("_")[5][1:]
    product_date = name.split("_")[2][:8].replace("-","")
    # producer = Producer("data")
    # producer.set_name_server_address(addr)
    # producer.start()
    data = {
        "tile": tile,  # 假设其他属性字段在结果中的索引为1
        "product_date": product_date,
        "cpjbdm": "1",
        "cpjdmc": "L1C",
        "imgtype_code": "S2",   
        "imgtype_name": "Sentinel-2",
        "cloud": float(read_cloud(zip_path)),
        "tif_path": tif_path,
        "geometry": tif_2_geo(tif_path)}
    json_path = os.path.join(json_dir,f"{name}.json")
    with open(json_path,"w") as json_file:
        json.dump(data,json_file,indent=4)
    redisr.rpush('raster_json_upload', f'"{json_path}"')
    # message_content = {"file_path":json_path}
    # message_body = json.dumps(message_content)
    # msg = Message(topic)
    # msg.set_tags("TagA")
    # msg.set_body(message_body)
    # _ = producer.send_sync(msg)
    # producer.shutdown()
    return 


def zip_2_jp2(zip_path, bands, save_dir):
    band_names = {
        "1":"T.*_B01\.jp2",
        "2":"T.*_B02\.jp2",
        "3":"T.*_B03\.jp2",
        "4":"T.*_B04\.jp2",
        "5":"T.*_B05\.jp2",
        "6":"T.*_B06\.jp2",
        "7":"T.*_B07\.jp2",
        "8A":"T.*_B8A\.jp2",
        "8":"T.*_B08\.jp2",
        "9":"T.*_B09\.jp2",
        "10":"T.*_B10\.jp2",
        "11":"T.*_B11\.jp2",
        "12":"T.*_B12\.jp2"}
    pathres = list(map(lambda x:".*/IMG_DATA/"+band_names[f"{x}"],bands))
    paths = []
    zipFile = zipfile.ZipFile(zip_path, "r")
    for pathre in pathres:
        for info in zipFile.infolist():
            if re.match(pathre,info.filename):
                path = zipFile.extract(info.filename, save_dir)
                paths.append(path)
    
    zipFile.close()
    return paths
    

def jp2s_2_tif(jp2_paths,tif_path,nodata=0):
    vrt_path = tif_path.replace('.tif','.vrt')
    options = gdal.BuildVRTOptions(
        resolution="highest",
        VRTNodata=nodata,
        resampleAlg=gdalconst.GRA_NearestNeighbour,
        separate=True)
    gdal.BuildVRT(vrt_path,jp2_paths,options=options)
    creationOptions = [
        "BIGTIFF=YES",
        "TILED=YES",
        "BLOCKXSIZE=1024",
        # "TFW=YES",
        "NUM_THREADS=ALL_CPUS"]
    options = gdal.TranslateOptions(
        format="GTiff",
        xRes=10,
        yRes=10,
        resampleAlg=gdalconst.GRA_NearestNeighbour,
        creationOptions=creationOptions)
    gdal.Translate(tif_path,vrt_path,options=options)
    os.remove(vrt_path)
    return


def zip_2_tif(zip_path,save_dir,bands=['2','3','4','5','6','7',"8A",'8','11','12']):
    zip_name = os.path.basename(zip_path)
    year_month = zip_name.split("_")[2][:6]
    date_dir = os.path.join(save_dir,year_month)
    if not os.path.exists(date_dir):os.makedirs(date_dir)
    tif_name = zip_name.replace('.SAFE', '').replace('.zip', '.tif')
    tif_path = os.path.join(date_dir, tif_name)
    # if os.path.exists(tif_path): return tif_path
    jp2_paths = zip_2_jp2(zip_path,bands,date_dir)
    jp2s_2_tif(jp2_paths,tif_path)
    unzip_dir = tif_path.replace(".tif",".SAFE")
    if os.path.exists(unzip_dir): shutil.rmtree(unzip_dir)
    return tif_path


def projection_2_spatil(tif_path):
    # 打开栅格数据集
    dataset = gdal.Open(tif_path)
    # 获取WKT格式的空间参考信息
    wkt = dataset.GetProjection()
    # 创建空间参考对象
    spatial_ref = osr.SpatialReference(wkt=wkt)
    # 检查空间参考是否存在
    if wkt == '': 
        raise "无空间参考信息"
    else:
        # 检查是否为地理坐标系
        if spatial_ref.IsGeographic():
            print("使用的是地理坐标系")
            # 获取EPSG代码
            epsg_code = spatial_ref.GetAttrValue('AUTHORITY', 1)
            # 检查是否是WGS84
            if epsg_code == '4326':
                print("该栅格使用的是WGS84坐标系。")
                # vrt_path = re.sub("\..*$",".vrt",tif_path)
                mem_path = '/vsimem/output.tif'
                gdal.Warp(mem_path, dataset, dstSRS='EPSG:4326')
                return mem_path
            else:
                # vrt_path = re.sub("\..*$",".vrt",tif_path)
                mem_path = '/vsimem/output.tif'
                # 使用gdal.Warp()进行重新投影
                gdal.Warp(mem_path, dataset, dstSRS='EPSG:4326')
                return mem_path
        # 检查是否为投影坐标系
        elif spatial_ref.IsProjected():
            print("使用的是投影坐标系")
            # 使用gdal.Warp()进行重新投影
            # vrt_path = re.sub("\..*$",".vrt",tif_path)
            mem_path = '/vsimem/output.tif'
            gdal.Warp(mem_path, dataset, dstSRS='EPSG:4326')
            return mem_path
        
        
def tif_2_geo(tif_path):
    # 打开栅格数据集
    raster_path = projection_2_spatil(tif_path)
    dataset = gdal.Open(raster_path,gdalconst.GA_Update)
    wkt_projection = dataset.GetProjection()
    band = dataset.GetRasterBand(1)
    nodata = band.GetNoDataValue()
    array = band.ReadAsArray()
    mask_band = band.GetMaskBand()
    mask_array = mask_band.ReadAsArray()
    band.WriteArray(mask_array)
    band = dataset.GetRasterBand(1)
    srs = osr.SpatialReference()
    srs.ImportFromWkt(wkt_projection)
    # 创建输出矢量数据源
    driver = ogr.GetDriverByName('Memory') # Memory
    mem_ds = driver.CreateDataSource('memory') # memory
    layer = mem_ds.CreateLayer('polygon_layer', srs, ogr.wkbPolygon)
    # 添加字段
    fieldDef = ogr.FieldDefn('bandvalue', ogr.OFTInteger)
    layer.CreateField(fieldDef)
    options = []
    options.append('8CONNECTED=8')  # 使用 8 连通区域生长算法
    gdal.Polygonize(
        band, 
        mask_band, 
        layer, 
        -1, 
        options=options, 
        callback=gdal.TermProgress_nocb
    )
    fearture = layer.GetFeature(0) 
    geometry = fearture.geometry()
    geometry_wkt = geometry.ExportToWkt()
    mem_ds.Destroy()
    del dataset
    return geometry_wkt


def main():
    with open('/data/jiabing/RS_CODE/Sentinel2download/config.json',"r") as file:
        params = json.load(file)
    json_dir = params['json_dir']
    tif_dir = params["tif_dir"]
    cargs = params["cargs"]
    rediscon = params["rediscon"]
    pool = redis.ConnectionPool(**rediscon)
    redisr = redis.Redis(connection_pool=pool)
    while True:
        time.sleep(10)
        _, element = redisr.blpop('raster_geometry',0)
        zip_path = element.decode('utf-8').replace('"','')
        # name = os.path.splitext(os.path.basename(zip_path))[0].replace(".SAFE","")
        # tif_path = name_2_count(cargs,name)
        # if tif_path is None: 
        tif_path = zip_2_tif(zip_path, tif_dir)
        send_data(json_dir,zip_path,tif_path,redisr)



if __name__ == "__main__":
    # main()
    # zip_path = r"/data/RS_Data/Image_Data/Sentinel2/1C/202404/S2A_MSIL1C_20240418T030501_N0510_R032_T50SMG_20240418T040141.zip"
    # tif_dir = r"/data/Yang/data/temp"
    # zip_2_tif(zip_path, tif_dir)
    with daemon.DaemonContext(
        stdout=open(r"/data/logfile/RasterGeometry_out.log","w"),
        stderr=open(r"/data/logfile/RasterGeometry_err.log","w")):
        main()